//
//  rtmp_streaming_buffer.cpp
//  RtmpSender
//
//  Created by LSQ on 2020/2/12.
//  Copyright © 2020 Agora. All rights reserved.
//

#include "rtmp_streaming_buffer.h"

#include <functional>

#include "utils/log/log.h"

namespace agora {
namespace rtc {

const char MODULE_NAME[] = "[RtmpStreamingBuffer]";

static const int kDefaultSortBufferMaxCount = 5;

// Update bitrate every 4 seconds if needed
static const int kDefaultCallBackIntervalMs = 4000;

// For the most common settings (15fps, 44.1k, mono), the buffer count in 1 second is:
// 15 + 44100 * 1 / 1024 = 58
static const int kDefaultUpdateIntervalMs = 800;

RtmpStreamingBuffer::RtmpStreamingBuffer(utils::worker_type worker, int max_buffer_size)
    : rtmp_send_worker_(worker),
      max_buffer_size_(max_buffer_size),
      current_interval_ms_(0),
      statistics_{0},
      buffer_observer_(utils::RtcAsyncCallback<IRtmpStreamingBufferObserver>::Create()) {
  if (!sample_timer_) {
    sample_timer_.reset(utils::minor_worker("AgRskSampleTick")
                            ->createTimer(std::bind(&RtmpStreamingBuffer::SampleTick, this),
                                          kDefaultUpdateIntervalMs));
  }
}

RtmpStreamingBuffer::~RtmpStreamingBuffer() {
  if (sample_timer_) {
    sample_timer_.reset();
  }
  CleanUp();
}

void RtmpStreamingBuffer::CleanUp() {
  current_interval_ms_ = 0;

  statistics_ = {0};
  std::lock_guard<std::mutex> guard_(list_mutex_);

  // clear frame_list_
  RemoveAllFrames();

  // clear sort_list
  for (auto frame : sort_list_) {
    if (frame.data != nullptr) {
      free(frame.data);
    }
  }
  sort_list_.clear();
  sample_list_.clear();
}

int RtmpStreamingBuffer::PushFrame(const RtmpMediaFrame& frame) {
  if (!frame.size) {
    return -1;
  }

  std::lock_guard<std::mutex> guard_(list_mutex_);
  if (frame.is_video_frame) {
    statistics_.push_video_frame_cnt_++;
    if (NeedLogStream()) {
      log(commons::LOG_DEBUG, "%s %s [%" PRId64 "] %s, pts: %u, size: %zu, key: %d, list size: %zu",
          MODULE_NAME, __FUNCTION__, statistics_.push_video_frame_cnt_, "video", frame.pts_ms,
          frame.size, frame.is_key_frame, frame_list_.size());
    }
  } else {
    statistics_.push_audio_frame_cnt_++;
    if (NeedLogStream()) {
      log(commons::LOG_DEBUG, "%s %s [%" PRId64 "] %s, pts: %u, size: %zu, key: %d, list size: %zu",
          MODULE_NAME, __FUNCTION__, statistics_.push_audio_frame_cnt_, "audio", frame.pts_ms,
          frame.size, frame.is_key_frame, frame_list_.size());
    }
  }
  sort_list_.push_back(frame);
  if (sort_list_.size() > kDefaultSortBufferMaxCount) {
    // sort
    sort_list_.sort();
    // drop frame if needed
    RemoveExpireFrame();
    // add frame to buffer
    RtmpMediaFrame first_frame = sort_list_.front();
    sort_list_.pop_front();
    frame_list_.push_back(first_frame);
  }
  return 0;
}

void RtmpStreamingBuffer::PopFirstFrame(RtmpMediaFrame& frame) {
  std::lock_guard<std::mutex> guard_(list_mutex_);
  // TODO(Haonong Yu): 2020/7/20 check empty
  frame = frame_list_.front();
  frame_list_.pop_front();

  if (frame.is_video_frame) {
    statistics_.pop_video_frame_cnt_++;
    if (NeedLogStream()) {
      log(commons::LOG_DEBUG, "%s %s [%" PRId64 "] %s, pts: %u, size: %zu, key: %d, list size: %zu",
          MODULE_NAME, __FUNCTION__, statistics_.pop_video_frame_cnt_, "video", frame.pts_ms,
          frame.size, frame.is_key_frame, frame_list_.size());
    }
  } else {
    statistics_.pop_audio_frame_cnt_++;
    if (NeedLogStream()) {
      log(commons::LOG_DEBUG, "%s %s [%" PRIu64 "] %s, pts: %u, size: %zu, key: %d, list size: %zu",
          MODULE_NAME, __FUNCTION__, statistics_.pop_audio_frame_cnt_, "audio", frame.pts_ms,
          frame.size, frame.is_key_frame, frame_list_.size());
    }
  }
}

int RtmpStreamingBuffer::FrameListSize() {
  std::lock_guard<std::mutex> guard_(list_mutex_);
  return static_cast<int>(frame_list_.size());
}

int RtmpStreamingBuffer::RegisterRtmpStreamingBufferObserver(
    IRtmpStreamingBufferObserver* observer) {
  API_LOGGER_MEMBER("observer: %p", observer);
  if (!observer) {
    return -ERR_INVALID_ARGUMENT;
  }
  return buffer_observer_->Register(observer);
}

int RtmpStreamingBuffer::UnregisterRtmpStreamingBufferObserver(
    IRtmpStreamingBufferObserver* observer) {
  API_LOGGER_MEMBER("observer: %p", observer);
  if (!observer) {
    return -ERR_INVALID_ARGUMENT;
  }
  return buffer_observer_->Unregister(observer);
}

RtmpStreamingBuffer::BufferStatistic RtmpStreamingBuffer::GetStatisticInfo() {
  API_LOGGER_MEMBER(nullptr);
  return statistics_;
}

int RtmpStreamingBuffer::RemoveTrailingPFrames() {
  API_LOGGER_MEMBER(nullptr);
  std::lock_guard<std::mutex> guard_(list_mutex_);
  return RemoveExpirePFrames(false);
}

void RtmpStreamingBuffer::RemoveExpireFrame() {
  if (frame_list_.size() < max_buffer_size_) {
    return;
  }

  int remove_pframe_count = RemoveExpirePFrames(true);
  if (remove_pframe_count > 0) {
    statistics_.skip_video_pframe_cnt_ += remove_pframe_count;
    return;
  }

  int remove_iframe_count = RemoveExpireIFrames();
  if (remove_iframe_count > 0) {
    statistics_.skip_video_iframe_cnt_ += remove_iframe_count;
    return;
  }

  // something is broken, clear frame list
  RemoveAllFrames();
}

// Remove the first group of P-frames
int RtmpStreamingBuffer::RemoveExpirePFrames(bool skipIFrame) {
  int remove_pframe_count = 0;
  for (auto iter = frame_list_.begin(); iter != frame_list_.end();) {
    RtmpMediaFrame frame = *iter;

    // Skip audio frames
    if (!frame.is_video_frame) {
      ++iter;
      continue;
    }

    // Skip I-frames before the first P-frame if needed
    if ((frame.is_key_frame && remove_pframe_count == 0)) {
      if (skipIFrame) {
        ++iter;
        continue;
      } else {
        break;
      }
    }

    if (frame.is_key_frame) {
      // Found the complete first group of P-frames
      break;
    }

    // Remove the P-frame
    if (frame.data != nullptr) {
      free(frame.data);
    }
    ++remove_pframe_count;
    log(commons::LOG_DEBUG, "%s %s pts: %u cnt: %d", MODULE_NAME, __FUNCTION__, frame.pts_ms,
        remove_pframe_count);
    iter = frame_list_.erase(iter);
  }
  return remove_pframe_count;
}

// Remove the first I-frame (which may include several NALUs) and trailing P-frames (just in case)
int RtmpStreamingBuffer::RemoveExpireIFrames() {
  int remove_iframe_count = 0;
  int remove_pframe_count = 0;
  uint32_t iframe_ts = 0;
  for (auto iter = frame_list_.begin(); iter != frame_list_.end();) {
    RtmpMediaFrame frame = *iter;

    // Skip audio frames
    if (!frame.is_video_frame) {
      ++iter;
      continue;
    }

    if (frame.is_key_frame && remove_iframe_count != 0 && frame.pts_ms != iframe_ts) {
      break;
    }

    // Remove the video frame
    if (frame.is_key_frame) {
      ++remove_iframe_count;
      iframe_ts = frame.pts_ms;
    } else {
      ++remove_pframe_count;
    }
    if (frame.data != nullptr) {
      free(frame.data);
    }
    log(commons::LOG_DEBUG, "%s %s pts: %u cnt: %d-%d", MODULE_NAME, __FUNCTION__, frame.pts_ms,
        remove_iframe_count, remove_pframe_count);
    iter = frame_list_.erase(iter);
  }
  return remove_iframe_count;
}

void RtmpStreamingBuffer::RemoveAllFrames() {
  log(commons::LOG_INFO, "%s %s count: %zu", MODULE_NAME, __FUNCTION__, frame_list_.size());
  for (auto frame : frame_list_) {
    if (frame.data != nullptr) {
      free(frame.data);
    }
  }
  frame_list_.clear();
}

RtmpStreamingBufferState RtmpStreamingBuffer::CurrentBufferState() {
  int last_buffer_count = 0;
  int increase_count = 0;
  int decrease_count = 0;

  auto iter = sample_list_.begin();
  last_buffer_count = *iter;
  while (++iter != sample_list_.end()) {
    int current_buffer_count = *iter;
    if (current_buffer_count > last_buffer_count) {
      increase_count++;
    } else if (current_buffer_count < last_buffer_count || current_buffer_count <= 0) {
      decrease_count++;
    }
    last_buffer_count = current_buffer_count;
  }

  // TODO(Haonong Yu): 2020/7/20 strategy modeling and no magic number
  static int kMaxIncreaseOrDeclineCount = kDefaultCallBackIntervalMs / kDefaultUpdateIntervalMs - 1;
  if (increase_count >= (kMaxIncreaseOrDeclineCount - 1)) {
    return kStreamingBufferIncreasing;
  } else if (decrease_count >= (kMaxIncreaseOrDeclineCount - 1)) {
    return kStreamingBufferDeclining;
  } else {
    return kStreamingBufferNone;
  }
}

void RtmpStreamingBuffer::SampleTick() {
  current_interval_ms_ += kDefaultUpdateIntervalMs;
  int frameListSize = FrameListSize();
  sample_list_.push_back(frameListSize);
  if (NeedLogStream()) {
    log(commons::LOG_INFO, "%s sample list bucket: %zu frame list size: %d", MODULE_NAME,
        sample_list_.size(), frameListSize);
  }
  if (current_interval_ms_ >= kDefaultCallBackIntervalMs) {
    RtmpStreamingBufferState state = CurrentBufferState();
    if (state != kStreamingBufferNone) {
      log(commons::LOG_DEBUG, "%s streaming buffer state: %d", MODULE_NAME, state);
      buffer_observer_->Post(
          LOCATION_HERE, [state](auto callback) { callback->OnRtmpStreamingBufferState(state); });
    }
    current_interval_ms_ = 0;
    sample_list_.clear();
  }
}

}  // namespace rtc
}  // namespace agora
